package io.daocloud.prometheustestdemo.service;

import io.github.bucket4j.Bandwidth;
import io.github.bucket4j.Bucket;
import io.github.bucket4j.Refill;
import io.github.bucket4j.redis.lettuce.cas.LettuceBasedProxyManager;
import lombok.extern.slf4j.Slf4j;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;
import org.springframework.beans.factory.annotation.Autowired;
import org.springframework.stereotype.Service;

import java.time.Duration;
import javax.annotation.PostConstruct;

/**
 * 限流服务
 */
@Service
@Slf4j
public class RateLimitService {
    
    private static final Logger log = LoggerFactory.getLogger(RateLimitService.class);
    
    @Autowired
    private LettuceBasedProxyManager proxyManager;
    
    private final int requestsPerSecond;
    
    public RateLimitService(@Autowired io.daocloud.prometheustestdemo.config.RateLimitConfig rateLimitConfig) {
        this.requestsPerSecond = rateLimitConfig.getRequestsPerSecond();
        log.info("RateLimitService initialized with {} requests per second", this.requestsPerSecond);
    }
    
    @PostConstruct
    public void init() {
        try {
            log.info("Testing Redis connection in RateLimitService...");
            // 测试创建一个bucket来验证Redis连接
            String testKey = "test_connection_key";
            getBucket(testKey);
            log.info("Redis connection test successful in RateLimitService");
        } catch (Exception e) {
            log.error("Redis connection test failed in RateLimitService: {}", e.getMessage());
            log.warn("Rate limiting will fall back to allowing all requests on Redis failures");
        }
    }
    
    /**
     * 检查是否允许请求
     * @param key 限流键（可以是IP地址、用户ID等）
     * @return true-允许请求，false-请求被限流
     */
    public boolean isAllowed(String key) {
        try {
            log.debug("Checking rate limit for key: {}", key);
            Bucket bucket = getBucket(key);
            boolean consumed = bucket.tryConsume(1);
            
            if (!consumed) {
                log.warn("Rate limit exceeded for key: {}", key);
            } else {
                log.debug("Rate limit check passed for key: {}", key);
            }
            
            return consumed;
        } catch (Exception e) {
            log.error("Rate limit check failed for key: {}, error: {}", key, e.getMessage());
            log.error("Exception details: ", e);
            // 如果Redis不可用，默认允许请求
            log.warn("Allowing request due to Redis failure");
            return true;
        }
    }
    
    /**
     * 获取或创建Bucket
     * @param key 限流键
     * @return Bucket实例
     */
    private Bucket getBucket(String key) {
        try {
            log.debug("Creating bucket for key: {}", key);
            // 每秒允许requestsPerSecond个请求
            Bandwidth limit = Bandwidth.classic(requestsPerSecond, Refill.intervally(requestsPerSecond, Duration.ofSeconds(1)));
            io.github.bucket4j.BucketConfiguration configuration = io.github.bucket4j.BucketConfiguration.builder()
                    .addLimit(limit)
                    .build();
            
            Bucket bucket = proxyManager.builder()
                    .build(key.getBytes(), configuration);
            log.debug("Bucket created successfully for key: {}", key);
            return bucket;
        } catch (Exception e) {
            log.error("Failed to create bucket for key: {}, error: {}", key, e.getMessage());
            throw e;
        }
    }
} 